package com.shujia.spark.sql

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.{DataFrame, Dataset, Row, SaveMode, SparkSession}

object Demo3DataFrameAPI {
  def main(args: Array[String]): Unit = {
    //1、创建spark sql环境
    val spark: SparkSession = SparkSession
      .builder()
      .master("local")
      .config("spark.sql.shuffle.partitions", 1)
      .appName("api")
      .getOrCreate()

    /**
     * 1、读取数据
     */
    val studentDF: DataFrame = spark
      .read
      .format("csv") //数据格式
      .option("sep", ",") //字段分隔符
      .schema("id STRING,name STRING,age INT, sex STRING,clazz STRING") //字段名和类型
      .load("data/students.txt")

    val scoreDF: DataFrame = spark
      .read
      .format("csv") //数据格式
      .option("sep", ",") //字段分隔符
      .schema("id STRING,cId STRING,score DOUBLE") //字段名和类型
      .load("data/score.txt")


    /**
     * 2、show：打印部分数据
     */
    studentDF.show()
    //指定查看的行数
    studentDF.show(100)
    //完全现在所有的数据
    studentDF.show(false)

    /**
     * 2、select:选择数据
     */
    //通过列名选择字段
    studentDF.select("name", "age").show()

    //导入隐式转换
    import spark.implicits._

    //使用列对象的方式，可以在select时对列进行处理
    studentDF.select($"name", $"age".+(1)).show()

    //可以省略逗号和括号
    studentDF.select($"name", $"age" + 20 as "age").show()


    /**
     * 3、where:过滤数据
     */
    //使用字符串表达式
    studentDF.where("sex = '男'").show()

    //使用列对象
    studentDF.where($"sex" === "女").show()


    /**
     * 4、groupBy  聚合
     */
    studentDF.groupBy($"clazz").count().show()

    //导入函数包
    import org.apache.spark.sql.functions._

    //使用agg进行聚合计算
    studentDF
      .groupBy($"clazz")
      .agg(count($"clazz") as "num", round(avg($"age"), 2) as "avgAge")
      .show()

    //直接使用where 代替having
    studentDF
      .groupBy($"clazz")
      .agg(count($"clazz") as "num", round(avg($"age"), 2) as "avgAge")
      .where($"num" > 90)
      .show()


    /**
     * 5、join:表关联
     *
     */

    //关联字段名不一样的写法
    //val joinDF: DataFrame = studentDF.join(scoreDF, $"id" === $"sid", "inner")

    //关联字段名一样时
    val joinDF: DataFrame = studentDF.join(scoreDF, "id")

    //加别名进行关联
    studentDF.as("a").join(scoreDF.as("b"), $"a.id" === $"b.id").show()

    joinDF.show()

    /**
     * orderBy: 排序
     */

    studentDF.orderBy($"age".desc).show()


    /**
     * limit
     */

    studentDF.orderBy($"age".desc).limit(10).show()


    /**
     * withColumn: 在原表上增加列
     */
    studentDF.withColumn("type", substring($"clazz", 1, 2)).show()

    /**
     * row_number:排名
     */
    //计算每个班级年龄最大的前十个学生
    studentDF
      .withColumn("r", row_number() over Window.partitionBy($"clazz").orderBy($"age".desc))
      .where($"r" <= 10)
      .show(1000)


    /**
     * 保存结果
     */
    val resultDF: DataFrame = studentDF
      .withColumn("r", row_number() over Window.partitionBy($"clazz").orderBy($"age".desc))
      .where($"r" <= 10)

    resultDF.write
      .format("csv")
      .option("sep", "\t")
      .mode(SaveMode.Overwrite)
      .save("data/top")

  }

}
